import torch
from torch import nn


class Module(nn.Module):
    def __init__(self):
        super(Module,self).__init__()
        self.models = nn.Sequential(nn.Conv2d(in_channels=3,out_channels=32,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=2),
                                    nn.Conv2d(in_channels=32,out_channels=32,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=(2,2)),
                                    nn.Conv2d(in_channels=32,out_channels=64,kernel_size=(5,5),padding=2),
                                    nn.MaxPool2d(kernel_size=(2,2)),
                                    nn.Flatten(),
                                    nn.Linear(1024,64),
                                    nn.Linear(64,10))

    def forward(self,x):
        x = self.models(x)
        return x

if __name__ == '__main__':
    model = Module()
    test1 = torch.ones([64,3,32,32])
    output = model(test1)
    print(output)
